"""
处理数据及特征选择的工具
"""
import numpy as np
from sklearn.preprocessing import OrdinalEncoder
from utils import data_loader
from utils import log
import os


def preprocess_data(df):
    """
    数据预处理和特征选择，保存日志
    :param df: 传入数据（df对象）
    :return: 返回处理过的df对象和特征列
    """
    log_path = os.path.join('../log', 'preprocessor')
    logger = log.setup_logger(log_path, 'preprocessor')

    # 增加特征列：HighRiskGroup、PromotionStagnationIndex
    logger.info('添加新特征: HighRiskGroup, PromotionStagnationIndex')

    # 定义4类高风险群体
    def is_high_risk(df):
        if (df['MonthlyIncome'] < 3000) and (df['JobLevel'] == 1) and (df['MaritalStatus'] == 'Single'):
            return 1
        elif (df['Age'] <= 29) and (df['YearsAtCompany'] <= 1) and (df['YearsSinceLastPromotion'] == 0):
            return 1
        elif df['JobRole'] in ['Laboratory Technician', 'Research Director'] and (df['OverTime'] == 'Yes'):
            return 1
        elif (df['Education'] in [1, 2]) and (df['MonthlyIncome'] <= 6000) and (df['JobSatisfaction'] <= 2):
            return 1
        else:
            return 0

    df['HighRiskGroup'] = df.apply(is_high_risk, axis=1)

    # 计算晋升停滞指数
    def calculate_stagnation_index(row):
        # 处理总工作年限为0的情况（刚入职员工）
        if row['TotalWorkingYears'] == 0:
            return -1

        # 处理晋升时间为0的情况（未晋升或刚晋升）
        years_since_promotion = row['YearsSinceLastPromotion']
        if years_since_promotion == 0:
            # 替代公式
            return row['Age'] / row['TotalWorkingYears']
        # 计算晋升停滞指数
        return (row['Age'] / row['TotalWorkingYears']) * (1 / years_since_promotion)

    df['StagnationIndex'] = df.apply(calculate_stagnation_index, axis=1)

    # 删除无用列及低价值特征
    logger.info('删除无用和代价值特征')
    df.dropna(axis=1, inplace=True)
    df = df.drop(['EmployeeNumber', 'Over18', 'StandardHours', 'TotalWorkingYears', 'YearsInCurrentRole',
                  'TrainingTimesLastYear', 'PercentSalaryHike', 'Department', 'NumCompaniesWorked', 'Gender'], axis=1)

    # 处理分类变量
    logger.info('处理分类变量')
    cat_cols = ['BusinessTravel', 'EducationField', 'JobRole', 'MaritalStatus', 'OverTime']
    num_cols = [col for col in df.columns if col not in cat_cols + ['Attrition']]

    # 分类变量编码
    logger.info('对分类变量进行编码')
    for col in cat_cols:
        # 检查是否存在缺失值，防止 fit_transform 报错
        if df[col].isnull().any():
            raise ValueError()

        oe = OrdinalEncoder(dtype=int)
        df[[col]] = oe.fit_transform(df[[col]])

    return df, num_cols + cat_cols


def save_data(df, save_path, file_name):
    """
    保存处理后的数据
    :param df: 处理后的数据（df对象）
    :param save_path: 数据保存路径
    :param file_name: 保存文件名
    :return:
    """
    # 配置日志
    log_path = os.path.join('../log', 'preprocessor')
    logger = log.setup_logger(log_path, 'preprocessor')
    # 保存处理后的数据
    data_loader.save_data(df, save_path, file_name)
    logger.info(f'数据保存成功，路径为：{save_path}{file_name}')


if __name__ == '__main__':
    data = data_loader.load_data('../data/', 'train.csv')
    preprocess_data(data)
    print(data.columns)
    save_data(data, '../data/', 'train_preprocessed.csv')
